from robot_arm import ETController, ROBOT_IP, P_REF
from aruco_detector import ArUcoDetector
from camera import Camera, ESC_KEY
from cam_calibration import JSON_PATH
import cv2
import json
import time
import os
import numpy as np


def degrees_to_radians(degrees):
    return np.deg2rad(degrees)


def convert_pose_to_robot_frame(tvecs, euler_angles, cali_trans):
    # Convert Euler angles from degrees to radians
    euler_angles_rad = degrees_to_radians(euler_angles)

    # Convert Euler angles to rotation matrix
    R_cam_to_marker, _ = cv2.Rodrigues(np.array(euler_angles_rad))
    t_cam_to_marker = np.array(tvecs).reshape((3, 1))

    # Create transformation matrix from camera to ArUco marker
    T_cam_to_marker = np.eye(4)
    T_cam_to_marker[:3, :3] = R_cam_to_marker
    T_cam_to_marker[:3, 3] = t_cam_to_marker.flatten()

    # Compute the transformation from robot end-effector to ArUco marker
    T_ee_to_marker = np.dot(cali_trans, T_cam_to_marker)

    # Extract the translation and rotation from the transformation matrix
    t_ee_to_marker = T_ee_to_marker[:3, 3]
    R_ee_to_marker = T_ee_to_marker[:3, :3]
    euler_angles_ee_to_marker, _ = cv2.Rodrigues(R_ee_to_marker)

    return t_ee_to_marker.flatten(), euler_angles_ee_to_marker.flatten()


if __name__ == "__main__":
    controller = ETController(ROBOT_IP)

    camera = Camera()
    if not camera.configure_pipeline():
        exit()
    if not camera.start_pipeline():
        exit()
    detector = ArUcoDetector()
    json_file_path = os.path.join(
        JSON_PATH, "calibration_result_20240603-103550.json"
    )  # JSON文件的路径
    if os.path.exists(json_file_path):
        with open(json_file_path, "r") as f:
            hand_cam_cali_result = json.load(f)
            cali_trans = np.array(hand_cam_cali_result["calibration_result"])
    else:
        print(f"JSON file {json_file_path} not found!")
        exit()

    start_time = time.time()
    frame_count = 0
    try:
        if controller.connect():
            # 获取机器人状态
            success_stat, result_stat, cmd_id_stat = controller.send_cmd(
                "getRobotState"
            )
            if success_stat:
                print("robot result:", result_stat)
            else:
                print("Failed to get robot state")

            while True:
                color_image = camera.get_frame()
                if color_image is None:
                    continue

                detected_image, detection_result = detector.detect_aruco(color_image)
                cv2.imshow("Color Viewer", detected_image)
                if detection_result.ids is not None:
                    # print(f"Detected IDs: {detection_result.ids}")
                    # print(f"Euler Angles (degrees): {detection_result.euler_angles}")
                    # print(f"Translation Vectors: {detection_result.tvecs}")

                    t_ee, euler_ee = convert_pose_to_robot_frame(
                        detection_result.tvecs,
                        detection_result.euler_angles,
                        cali_trans,
                    )

                    # print(f"Translated Tvecs to EE: {t_ee}")
                    # print(f"Translated Euler to EE: {euler_ee}")

                frame_count += 1
                elapsed_time = time.time() - start_time
                if elapsed_time > 1.0:
                    fps = frame_count / elapsed_time
                    print(f"FPS: {fps:.2f}")
                    start_time = time.time()
                    frame_count = 0

                key = cv2.waitKey(1)
                if key == ESC_KEY:
                    break
                elif key == ord("w"):
                    print("*************************")
                    t_ee[2] -= 100
                    P_tar = []
                    P_tar.extend(t_ee.tolist())
                    P_tar.extend(euler_ee.tolist())
                    suc_inverse, result_inverse, id = controller.send_cmd(
                        "inverseKinematic",
                        {
                            "targetPose": P_tar,
                            "referencePos": P_REF,
                            "unit_type": 0,
                        },
                    )
                    if suc_inverse:
                        print("robot result:", result_inverse)
                        suc_move, result, id = controller.send_cmd(
                            "moveByLine",
                            {
                                "targetPos": result_inverse,
                                "speed_type": 0,
                                "speed": 200,
                                "cond_type": 0,
                                "cond_num": 7,
                                "cond_value": 1,
                            },
                        )
                        time.sleep(5)
                        print("robot arm will raise up")
                        P_tar[2] -= 100
                        suc, result_inverse, id = controller.send_cmd(
                            "inverseKinematic",
                            {
                                "targetPose": P_tar,
                                "referencePos": P_REF,
                                "unit_type": 0,
                            },
                        )
                        if suc:
                            print("robot result:", result_inverse)
                            suc, result, id = controller.send_cmd(
                                "moveByLine",
                                {
                                    "targetPos": result_inverse,
                                    "speed_type": 0,
                                    "speed": 200,
                                    "cond_type": 0,
                                    "cond_num": 7,
                                    "cond_value": 1,
                                },
                            )
                    else:
                        print("Failed to get inverse kinematic solution")

    except KeyboardInterrupt:
        pass
    finally:
        camera.stop_pipeline()
        cv2.destroyAllWindows()
